﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace LPLib
{
    public enum MonomerType { Nucleotide, AminoAcid }
    public enum RepresentationType { Pure, Mixed, PMF }

    public static class SymbolSets
    {
        public static char[] PureNucleotides = new char[] { 'A', 'G', 'T', 'C', '-' };
        public static char[] MixedNucleotides = new char[] { ' ', 'A', 'G', 'R', 'T', 'W', 'K', 'D', 'C', 'M', 'S', 'V', 'Y', 'H', 'B', 'N', '-' };
        public static char[] AminoAcids = new char[] { 'A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y', '*', '-'};
    }

    public abstract class MonomerEncoding<T>
    {
        protected System.Type basicType;
        protected RepresentationType repType;
        public RepresentationType RepType
        {
            get { return repType; }
        }
        public abstract bool IsValid(T monomer);
        public abstract bool IsGap(T monomer);
        public abstract T Complement(T monomer);
        public abstract T FromChar(char c);
        public abstract char ToChar(T monomer);

        public abstract string ToString(T[] sequence);

        public abstract T[] FromString(string sequence);

        public abstract T Gap();
    }

    public abstract class NucleotideEncoding<T> : MonomerEncoding<T>
    {
        public static double LogPrior = -1.38629436; // = Log(4)
    }

    public class AminoAcidChar : MonomerEncoding<char>
    {
        public override char Complement(char monomer)
        {
            throw new NotImplementedException();
        }

        public override char FromChar(char c)
        {
            return c;
        }

        public override bool IsGap(char monomer)
        {
            return monomer == '-';
        }

        public override bool IsValid(char monomer)
        {
            return SymbolSets.AminoAcids.Contains(monomer) && (monomer != '*') && (monomer != '-');
        }

        public override char ToChar(char monomer)
        {
            return monomer;
        }

        public override string ToString(char[] sequence)
        {
            StringBuilder sb = new StringBuilder();
            foreach (char c in sequence)
            {
                sb.Append(c);
            }
            return sb.ToString();
        }

        public override char[] FromString(string sequence)
        {
            char[] seq = new char[sequence.Length];
            for (int i = 0; i < seq.Length; i++)
            {
                seq[i] = sequence[i];
            }
            return seq;
        }

        public override char Gap()
        {
            return '-';
        }
    }

    public class AminoAcidPMF : MonomerEncoding<double[]>
    {
        private static Dictionary<char, int> encoder;
        private int dim = SymbolSets.AminoAcids.Length;
        public int Dim
        {
            get { return dim; }
        }

        static AminoAcidPMF()
        {
            encoder = new Dictionary<char, int>();
            for (int i= 0; i < SymbolSets.AminoAcids.Length; i++)
            {
                encoder.Add(SymbolSets.AminoAcids[i], i);
            }
        }

        public override double[] Complement(double[] monomer)
        {
            throw new NotImplementedException();
        }

        public override double[] FromChar(char c)
        {
            double[] pmf = new double[dim];
            if (encoder.ContainsKey(c))
            {
                pmf[encoder[c]] = 1;
            }
            return pmf;
        }

        public override double[][] FromString(string sequence)
        {
            double[][] seq = new double[sequence.Length][];
            for (int i = 0; i < sequence.Length; i++ )
            {
                seq[i] = FromChar(sequence[i]);
            }
            return seq;
        }

        public override bool IsGap(double[] a)
        {
            return a[21] == 1;
        }

        public override bool IsValid(double[] a)
        {
            if (a.Length != Dim)
                return false;
            for (int i = 0; i < Dim; i++)
            {
                if (a[i] < 0)
                    return false;
            }

            if (a.Sum() != 1)
                return false;

            return true; 
        }

        public override string ToString(double[][] sequence)
        {
            StringBuilder sb = new StringBuilder();
            for (int i = 0; i < sequence.Length; i++)
            {
                sb.Append(ToChar(sequence[i]));
            }
            return sb.ToString();
        }
        public override char ToChar(double[] a)
        {
            if (a.Length != Dim)
                throw new ArgumentException("The length of the argument is incorrect.");

            double max = a.Max();
            int i = 0;
            for (; i < a.Length; i++ )
            {
                if (a[i] == max)
                {
                    return SymbolSets.AminoAcids[i];
                }
            }

            return ' ';
        }

        public override double[] Gap()
        {
            double[] gap = new double[22];
            gap[21] = 1;
            return gap;
        }
    }

    /// <summary>
    /// Encoding of the pure nucleotides A, G, T, C by the integers 0, 1, 2, 3.
    /// </summary>
    public class PureNucleotideInt : NucleotideEncoding<int>
    {
        // A : 0
        // G : 1
        // T : 2
        // C : 3
        // - : 4
        // gap included for convenience
        // n1, n2 related by transition iff n1/2 = n2/2 in integer arithmetic
        // n1, n2 complementary iff n1 % 2 = n2 % 2

        private static char[] symbols = SymbolSets.PureNucleotides;
        private static Dictionary<char, int> encoder;
        public static Dictionary<char, int> Encoder
        {
            get { return PureNucleotideInt.encoder; }
        }
        private static int[] complement = new int[] { 2, 3, 0, 1, 4 };
        private static int[] transition = new int[] { 1, 0, 3, 2, 4 };
        private static int[] transversion2 = new int[] { 3, 2, 1, 0, 4 };
        private static Dictionary<int, char> decoder;

        static PureNucleotideInt()
        {
            encoder = new Dictionary<char, int>();
            decoder = new Dictionary<int, char>();
            for (int i = 0; i < symbols.Length; i++)
            {
                encoder.Add(symbols[i], i);
                decoder.Add(i, symbols[i]);
            }
            decoder.Add(-1, 'x');
        }
        /// <summary>
        /// Gets the Watson-Crick complement of the argument
        /// </summary>
        /// <param name="n">An integer representing a pure nucleotide.</param>
        /// <returns>The integer representing the Watson-Crick complement of the argument. Returns -1 if the argument is illegitimate.</returns>
        public override int Complement(int n)
        {
            if (n < 0 || n > 4)
                return -1;
            else
                return complement[n];
        }

        /// <summary>
        /// Determines whether a given variable of the appropriate type represents a gap.
        /// </summary>
        /// <param name="n">An integer representing a pure nucleotide.</param>
        /// <returns>True if and only if n represets a gap.</returns>
        public override bool IsGap(int n)
        {
            return n == 4;
        }

        /// <summary>
        /// Determines whether a given integer represents a pure nucleotide.
        /// </summary>
        /// <param name="n">An integer representing a pure nucleotide.</param>
        /// <returns>True if and only if n represets a pure nucleotide.</returns>
        public override bool IsValid(int n)
        {
            return (n > -1 && n < 5);
        }

        /// <summary>
        /// Gets the nucleotide that is the transition partner of the argument.
        /// </summary>
        /// <param name="n">A nucleotide in integer representation.</param>
        /// <returns>The nucleotide that represents a transition mutation from the argument nucleotide, or -1 if the argument is illegitimate.</returns>
        public static int Transition(int n)
        {
            if (n < 0 || n > 4)
                return -1;
            else
                return transition[n];
        }

        /// <summary>
        /// Gets the nucleotide that is the transversion2 (not-complement) partner of the argument.
        /// </summary>
        /// <param name="n">A nucleotide in integer representation.</param>
        /// <returns>The nucleotide that represents a transversion-not-complement mutation from the argument nucleotide, or -1 if the argument is illegitimate.</returns>
        public static int Transversion2(int n)
        {
            if (n < 0 || n > 4)
                return -1;
            else
                return transversion2[n];
        }

        public override int FromChar(char symbol)
        {
            if (symbols.Contains(symbol))
                return encoder[symbol];
            else
                return -1;
        }

        public override char ToChar(int monomer)
        {
            return TypeConverter.ToChar(monomer);
        }

        public override string ToString(int[] sequence)
        {
            StringBuilder s = new StringBuilder();
            foreach (int n in sequence)
            {
                s.Append(decoder[n]);

            }
            return s.ToString();
        }

        public override int[] FromString(string sequence)
        {
            int[] array = new int[sequence.Length];
            for (int i= 0; i < sequence.Length; i++)
            {
                if (encoder.ContainsKey(sequence[i]))
                    array[i] = encoder[sequence[i]];
                else
                    array[i] = -1;
            }
            return array;
        }

        public override int Gap()
        {
            return 4;
        }
    }

    /// <summary>
    /// Provides support for for nucleotides encoded by bytes.
    /// </summary>
    public class MixedNucleotideByte : NucleotideEncoding<byte>
    {
        private static char[] symbols = SymbolSets.MixedNucleotides;
        private static Dictionary<char, byte> encoder;
        public static Dictionary<char, byte> Encoder
        {
            get { return MixedNucleotideByte.encoder; }
        }
        private static Dictionary<byte, char> decoder;
        private static byte[] complement;
        private static int[] degeneracy;
        private static double[,] innerProduct;

        public const int NSymbols = 17;

        static MixedNucleotideByte()
        {
            complement = new byte[] { 0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15, 16 };
            degeneracy = new int[] { 0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 1 };
            encoder = new Dictionary<char, byte>();
            decoder = new Dictionary<byte, char>();
            for (byte b = 0; b < NSymbols; b++)
            {
                encoder.Add(symbols[b], b);
                decoder.Add(b, symbols[b]);
            }
            computeInnerProducts();
        }

        private static void computeInnerProducts()
        {
            innerProduct = new double[NSymbols, NSymbols];
            for (int i = 0; i < NSymbols; i++)
            {
                for (int j = 0; j < NSymbols; j++)
                {
                    ByteBits b3 = new ByteBits((byte)(i & j));
                    innerProduct[i, j] = b3.NSet() / ((double)degeneracy[i] * degeneracy[j]);
                }
            }
        }

        /// <summary>
        /// Reports whether a given byte legitimately encodes a nucleotide.
        /// </summary>
        /// <param name="n">The nucleotide to be tested.</param>
        /// <returns>True if n encodes a nucleotide. False otherwise.</returns>
        public override bool IsValid(byte n)
        {
            return (n < 17);
        }

        /// <summary>
        /// Computes the complement of a byte-encoded nucleotide.
        /// </summary>
        /// <param name="n">The byte-encoded nucleotide to be complemenented.</param>
        /// <returns>The Watson-Crick complement of n. If n is not valid, returns 0. </returns>
        public override byte Complement(byte n)
        {
            if (IsValid(n))
            {
                return complement[n];
            }
            else
            {
                return 0;
            }
        }

        /// <summary>
        /// Computes the IUPAC degeneracy of a byte-encoded nucleotide symbol.
        /// </summary>
        /// <param name="n">A byte-encoded nucleotide symbol.</param>
        /// <returns>The number of pure nucleotides consistent with n.</returns>
        public static int Degeneracy(byte n)
        {
            if (n < 17)
            {
                return degeneracy[n];
            }
            else
            {
                return 0;
            }
        }

        /// <summary>
        /// Determines whether a nucleotide symbol encodes a gap.
        /// </summary>
        /// <param name="n">A byte-encoded nucleotide symbol.</param>
        /// <returns>True if n encodes a gap. False otherwise, including if n is illegitimate..</returns>
        public override bool IsGap(byte n)
        {
            return n == 16;
        }

        /// <summary>
        /// Determines whether two byte-encoded nucleotide symbols are consistent.
        /// </summary>
        /// <param name="n1">A byte-encoded nucleotide symbol.</param>
        /// <param name="n2">A byte-encoded nucleotide symbol.</param>
        /// <returns>True if there is a nucleotide that is consistent with both n1 and n2. False otherwise. The result is meanngless if either n1 or n2 is illegitimate.</returns>
        public static bool AreConsistent(byte n1, byte n2)
        {
            return (n1 & n2) > 0;
        }

        public static double InnerProduct(byte n1, byte n2)
        {
            if (n1 > 17 || n2 > 17) return double.NaN;

            return innerProduct[n1, n2];
        }

        public static bool RepresentsPureState(byte n)
        {
            return (Degeneracy(n) == 1);
        }

        public static int ToPureState(byte n)
        {
            switch (n)
            {
                case 1:
                    return 0;
                case 2:
                    return 1;
                case 4:
                    return 2;
                case 8:
                    return 3;
                default:
                    return -1;
            }
        }

        public override byte FromChar(char symbol)
        {
            if (symbols.Contains(symbol))
                return encoder[symbol];
            else
                return 0;
        }

        public override char ToChar(byte monomer)
        {
            return TypeConverter.ToChar(monomer);
        }

        public override byte[] FromString(string sequence)
        {
            byte[] array = new byte[sequence.Length];
            for (int i = 0; i < sequence.Length; i++)
            {
                if (encoder.ContainsKey(sequence[i]))
                {
                    array[i] = encoder[sequence[i]];
                }
                else
                {
                    array[i] = 0;
                }
            }
            return array;
        }

        public override string ToString(byte[] sequence)
        {
            StringBuilder s = new StringBuilder();
            foreach (byte b in sequence)
            {
                s.Append(decoder[b]);
            }
            return s.ToString();
        }

        public override byte Gap()
        {
            return 16;
        }
    }

    /// <summary>
    /// Provides support for for nucleotides encoded by bytes.
    /// </summary>
    public class MixedNucleotideChar : NucleotideEncoding<char>
    {
        private static char[] symbols = SymbolSets.MixedNucleotides;
        private static char[] symbolComplements = SymbolSets.MixedNucleotides;
        private static Dictionary<char, char> complement;
        public const int NSymbols = 17;

        static MixedNucleotideChar()
        {
            symbolComplements = new char[] { ' ', 'T', 'C', 'Y', 'A', 'W', 'M', 'H', 'G', 'K', 'S', 'B', 'R', 'D', 'V', 'N', '-' };
            complement = new Dictionary<char, char>();
            for (int i = 0; i < symbols.Length; i++)
            {
                complement.Add(symbols[i], symbolComplements[i]);
            }
        }

        /// <summary>
        /// Reports whether a given byte legitimately encodes a nucleotide.
        /// </summary>
        /// <param name="n">The nucleotide to be tested.</param>
        /// <returns>True if n encodes a nucleotide. False otherwise.</returns>
        public override bool IsValid(char c)
        {
            return symbols.Contains(c);
        }

        /// <summary>
        /// Computes the complement of a byte-encoded nucleotide.
        /// </summary>
        /// <param name="n">The byte-encoded nucleotide to be complemenented.</param>
        /// <returns>The Watson-Crick complement of n. If n is not valid, returns 0. </returns>
        public override char Complement(char c)
        {
            if (IsValid(c))
            {
                return complement[c];
            }
            else
            {
                return ' ';
            }
        }

        /// <summary>
        /// Determines whether a nucleotide symbol encodes a gap.
        /// </summary>
        /// <param name="n">A byte-encoded nucleotide symbol.</param>
        /// <returns>True if n encodes a gap. False otherwise, including if n is illegitimate..</returns>
        public override bool IsGap(char c)
        {
            return c == '-';
        }

        public override char FromChar(char symbol)
        {
            return symbol;
        }

        public override char ToChar(char monomer)
        {
            return monomer;
        }

        public override string ToString(char[] sequence)
        {
            throw new NotImplementedException();
        }

        public override char[] FromString(string sequence)
        {
            throw new NotImplementedException();
        }

        public override char Gap()
        {
            return '-';
        }
    }

    /// <summary>
    /// Provides support for nucleotides as probability mass functions (PMF), that is, double[].
    /// </summary>
    public class NucleotidePMF : NucleotideEncoding<double[]>
    {
        public static int Dim = 5;
        public static MixedNucleotideByte mnbEncoding = new MixedNucleotideByte();
        /// <summary>
        /// Constructor. Doesn't do anything.
        /// </summary>
        public NucleotidePMF()
        {

        }

        /// <summary>
        /// Determines whether a double[] encodes a legitimate nucleotide (PMF).
        /// </summary>
        /// <param name="n">A double array.</param>
        /// <returns>True if n encodes a legitimate nucleotidet PMF; false otherwise.</returns>
        public override bool IsValid(double[] n)
        {
            if (n.Length != Dim)
                return false;
            for (int i = 0; i < Dim; i++)
            {
                if (n[i] < 0)
                    return false;
            }

            if (n.Sum() != 1)
                return false;

            return true;
        }

        /// <summary>
        /// The complement of a double[] variable treated as a nucleotide PMF.
        /// </summary>
        /// <param name="n">A double array.</param>
        /// <returns>The Watson-Crick complement of n. Return value is meaningless if n is illegitimate.</returns>
        public override double[] Complement(double[] n)
        {
            double[] comp = new double[Dim];
            for (int i = 0; i < Dim; i++)
            {
                comp[i] = n[(i + 2) % 4];
            }
            return comp;
        }

        /// <summary>
        /// Determines whether a double array represents a gap state.
        /// </summary>
        /// <param name="n">A double[] variable representing a nucleotide symbol.</param>
        /// <returns>True if the argument represents a pure gap state, False otherwise. Throws an exception of n.Length is less than 5</returns>
        public override bool IsGap(double[] n)
        {
            return n[4] == 1;
        }

        public override double[] FromChar(char c)
        {
            double[] pmf = new double[NucleotidePMF.Dim];
            byte b = mnbEncoding.FromChar(c);
            return TypeConverter.ToPMF(b);
        }

        public override char ToChar(double[] monomer)
        {
            return TypeConverter.ToChar(monomer);
        }

        public override double[][] FromString(string sequence)
        {
            double[][] seq = new double[sequence.Length][];
            for (int i= 0; i < sequence.Length; i++)
            {
                seq[i] = FromChar(sequence[i]);
            }
            return seq;
        }

        public override string ToString(double[][] sequence)
        {
            throw new NotImplementedException();
        }

        public static double IDFlux(double[] n1, double[] n2)
        {
            if (n1.Length != Dim || n2.Length != Dim)
                throw new ArgumentException();

            return n1[0] * n2[0] + n1[1] * n2[1] + n1[2] * n2[2] + n1[3] * n2[3] + n1[4] * n2[4];
        }

        public static double TransitionFlux(double[] n1, double[] n2)
        {
            if (n1.Length != Dim || n2.Length != Dim)
                throw new ArgumentException();

            return n1[0] * n2[1] + n1[1] * n2[0] + n1[2] * n2[3] + n1[3] * n2[2]; 
        }

        public static double TransversionFlux(double[] n1, double[] n2)
        {
            if (n1.Length != Dim || n2.Length != Dim)
                throw new ArgumentException();

            return n1[0] * (n2[2] + n2[3]) + n1[1] * (n2[2] + n2[3]) + n1[2] * (n2[0] + n2[1]) + n1[3] * (n2[0] + n2[1]);
        }

        public static double GapFlux(double[] n1, double[] n2)
        {
            if (n1.Length != Dim || n2.Length != Dim)
                throw new ArgumentException();

            return (1 - n1[4]) * n2[4] + (1 - n2[4]) * n1[4];
        }
        public override double[] Gap()
        {
            return new double[] { 0, 0, 0, 0, 1 };
        }

        public static double[] Normalize(double[] likelihood)
        {
            double[] pmf = new double[likelihood.Length];
            double sum = likelihood.Sum();
            if (sum == 0)
            {
                return likelihood;
            }

            for (int i= 0; i < likelihood.Length; i++)
            {
                pmf[i] = likelihood[i] / sum;
            }
            return pmf;
        }

        public static double[][] Normalize(double[][] likelihoodSequence)
        {
            double[][] pmfSequence = new double[likelihoodSequence.Length][];
            for (int i= 0; i < likelihoodSequence.Length; i++)
            {
                pmfSequence[i] = Normalize(likelihoodSequence[i]);
            }
            return pmfSequence;
        }
    }
}
